from rest_framework.views import APIView
from rest_framework.request import Request
from rest_framework.serializers import ModelSerializer
from utils.response import APIResponse
from pathlib import Path
from studio_celery.aicv.tasks import aicv
from celery.result import AsyncResult
from studio_celery.main import app
from django.conf import settings
from .models import AicvModel
from django.http import HttpResponse, FileResponse
from django.core.files.storage import FileSystemStorage


class AicvSerializer(ModelSerializer):
    class Meta:
        model = AicvModel
        fields = '__all__'


class AicvAPIView(APIView):
    queryset = None
    serializer_class = None

    def get(self, request: Request):
        data = AicvModel.objects.order_by('-ctime').all()
        data = AicvSerializer(instance=data, many=True)
        return APIResponse(data=data.data)

    def post(self, request: Request):
        file = request.FILES['file']
        filePath = settings.MEDIA_ROOT / file.name
        with filePath.open('wb+') as fp:
            for chunk in file.chunks():
                fp.write(chunk)

        result = aicv.delay(str(filePath))
        AicvModel.objects.create(
            task_id=result.id,
            result=file.name
        ).save()
        return APIResponse(code=200, data={'task_id': result.id}, msg='提交成功')

    def delete(self, request: Request):
        task_id = request.data.get('task_id')
        AicvModel.objects.filter(task_id=task_id).delete()
        result = AsyncResult(id=task_id, app=app)
        file_stream = Path(result.get())
        file_stream.unlink()
        result.forget()
        return APIResponse(code=200, msg='删除成功')


class AicvStatusAPIView(APIView):
    queryset = None
    serializer_class = None

    def get(self, request: Request):
        task_id = request.query_params.get('task_id')
        if not task_id:
            return APIResponse(code=200, msg='参数错误')
        async_result = AsyncResult(id=task_id, app=app)
        if async_result.successful():
            # 获取执行结果
            result = async_result.get()
            # .forget() 将结果删除，执行完成，结果不会自动删除
            return APIResponse(code=200, msg='OK', data={'result': result})
        elif async_result.failed():
            return APIResponse(code=200, msg='执行失败', data={'current': 1, 'total': 1})
        elif async_result.status == 'ERROR':
            return APIResponse(code=200, msg='执行失败', data={'current': 1, 'total': 1})
        elif async_result.status == 'PENDING':
            return APIResponse(code=200, msg='PENDING', data={'current': 1, 'total': 1})
        elif async_result.status == 'RETRY':
            return APIResponse(code=200, msg='任务异常正在重试', data={'current': 1, 'total': 1})
        elif async_result.status == 'STARTED':
            return APIResponse(code=200, msg='任务已经开始被执行', data={'current': 1, 'total': 1})
        elif async_result.status == 'PROGRESS':
            return APIResponse(code=200, msg='PROGRESS', data={
                'current': async_result.info['current'],
                'total': async_result.info['total']})
        return APIResponse(code=500, msg='error', data={'current': 1, 'total':1})

    def post(self, request: Request):
        task_id = request.data.get('task_id')
        async_result = AsyncResult(id=task_id, app=app)
        if async_result.successful():
            # 获取执行结果
            result = async_result.get()
            # file_stream = Path(r'/data/ai_studio/media/400mb.zip')
            file_stream = Path(result)
            response = FileResponse(file_stream.open('rb'))
            response['Content-Type'] = 'application/octet-stream'
            response['Content-Disposition'] = 'attachment;filename="{0}"'.format(file_stream.name)
            return response
        else:
            return APIResponse(code=400, msg='当前任务不可下载')
